[https://nvbugs/6480621][test] Revert to 60-second KV transfer timeout for GB300 DeepSeek V4 Pro disaggregated perf-sanity - #17137

tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py (1)

132-139: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Annotate and document the shared launch-script interface.

precheck_prefix_lines is called by both submit modules. Add precise parameter and return annotations. Add Google-style Args and Returns sections that document llm_models_root and the generated export lines.

As per coding guidelines, “Annotate every function” and “Prefer docstrings for external interfaces, use Google-style docstrings, document public function arguments.”

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py`
around lines 132 - 139, Update precheck_prefix_lines with precise type
annotations for every parameter and its return value, and add a Google-style
docstring documenting all arguments—especially llm_models_root—and that the
function returns generated export lines. Keep the shared launch-script interface
behavior unchanged.

Source: Coding guidelines

tests/unittest/others/test_cache_transceiver_precheck_run.py (1)

380-397: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Clarify the two event lists.

_ctx_finish_runner binds its own list to free_events, and the test then overrides _free_all to append to the local events. The assertion free_events == [] therefore checks the discarded list. The test is correct, but the names invert the reader's expectation. Take the runner list as _ or drop the override and assert on the runner list only.

♻️ Proposed simplification

-def test_ctx_finish_wave_frees_only_after_block_all_returns_every_request(monkeypatch):
-    events = []
-
-    def check_status(at_least_request_num):
-        events.append(("block_all", at_least_request_num))
-        return [101, 102], []
-
-    runner, free_events = _ctx_finish_runner(monkeypatch, check_status)
+def test_ctx_finish_wave_frees_only_after_block_all_returns_every_request(monkeypatch):
+    calls = []
+
+    def check_status(at_least_request_num):
+        calls.append(("block_all", at_least_request_num))
+        return [101, 102], []
+
+    runner, events = _ctx_finish_runner(monkeypatch, check_status)
     reqs = {
     }
-    runner._free_all = lambda owned: events.append(("free", sorted(owned)))
 
     runner.ctx_finish_wave(reqs)
 
-    assert events == [("block_all", None), ("free", [0, 1])]
-    assert free_events == []
+    assert calls == [("block_all", None)]
+    assert events == [("free", [0, 1])]

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/others/test_cache_transceiver_precheck_run.py` around lines
380 - 397, Clarify the event-list usage in
test_ctx_finish_wave_frees_only_after_block_all_returns_every_request by
discarding the unused free_events value from _ctx_finish_runner and asserting
the _free_all callback’s local events directly, or otherwise assert only the
runner-owned list without checking the discarded list.

tests/unittest/disaggregated/test_transceiver_bounded_polling.py (1)

421-428: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

This test does not reach the in-loop sibling check.

wait_complete calls has_failed() before the KV loop, so the pre-existing TaskStatus.ERROR on failed_task returns WaitResult.FAILED immediately. The assertions pending_task.wait_calls == [] confirm that. The test therefore duplicates test_tx_session_blocking_wait_treats_task_failure_as_terminal and leaves the sibling recheck at lines 1368-1371 of tensorrt_llm/_torch/disaggregation/native/transfer.py uncovered.

To cover that path, make the sibling fail during the first wait slice.

💚 Proposed test change to exercise the sibling recheck

 def test_tx_session_blocking_wait_detects_failed_sibling_behind_pending_task() -> None:
     pending_task = _FakeTask(TaskStatus.TRANSFERRING, wait_result=False)
-    failed_task = _FakeTask(TaskStatus.ERROR)
+    failed_task = _FakeTask(TaskStatus.TRANSFERRING, wait_result=False)
     session = _make_tx_session([pending_task, failed_task])
+    wait = pending_task.wait
+
+    def fail_sibling_during_wait(timeout: Optional[float] = None) -> bool:
+        result = wait(timeout)
+        failed_task.status = TaskStatus.ERROR
+        return result
+
+    pending_task.wait = fail_sibling_during_wait
 
     assert session.wait_complete(blocking=True) == WaitResult.FAILED
-    assert pending_task.wait_calls == []
+    assert pending_task.wait_calls == [0.25]
     assert failed_task.wait_calls == []

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/disaggregated/test_transceiver_bounded_polling.py` around
lines 421 - 428, Update
test_tx_session_blocking_wait_detects_failed_sibling_behind_pending_task so
failed_task starts in a non-error state and transitions to TaskStatus.ERROR
during pending_task’s first wait slice, allowing wait_complete(blocking=True) to
reach and validate the in-loop sibling failure recheck. Preserve the assertions
that the result is WaitResult.FAILED and both tasks’ wait-call behavior remains
correct.

tests/scripts/perf-sanity/disaggregated/gb300_deepseek-v4-pro-fp4_8k1k_con180_ctx3_dep4_gen1_dep32_eplb384_mtp3_ccb-NIXL.yaml (1)

70-70: 🩺 Stability & Availability | 🔵 Trivial

Note the reduced receive deadline for the high-concurrency workload.

kv_transfer_timeout_ms feeds rx_timeout_s in KvCacheTransceiverV2.__init__ (tensorrt_llm/_torch/disaggregation/transceiver.py line 108). The 10x reduction to 60 s tightens the receive deadline. The PR description states that the targeted run does not establish resolution of the concurrency-1760, 8-CTX-worker case. Under that load a cold NIXL link can exceed 60 s and the request then fails instead of completing late. Track a follow-up run at the original concurrency before this config is used as the perf-sanity baseline.

Also applies to: 101-101

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@tests/scripts/perf-sanity/disaggregated/gb300_deepseek-v4-pro-fp4_8k1k_con180_ctx3_dep4_gen1_dep32_eplb384_mtp3_ccb-NIXL.yaml`
at line 70, Restore kv_transfer_timeout_ms to its previous value for this
high-concurrency perf-sanity configuration, rather than using the reduced 60000
ms receive deadline. Apply the same correction to the additionally referenced
occurrence and retain the original timeout until the concurrency-1760,
8-CTX-worker follow-up run validates a shorter deadline.

tensorrt_llm/_torch/disaggregation/native/transfer.py (1)

1374-1390: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win

Remove the unreachable WaitResult.TIMEOUT handling. No wait_complete implementation returns WaitResult.TIMEOUT; remove the branch and timed_out plumbing from _ctx_consensus_outcome.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/disaggregation/native/transfer.py` around lines 1374 -
1390, The _ctx_consensus_outcome flow still contains obsolete timeout handling.
Remove the unreachable WaitResult.TIMEOUT branch and all timed_out plumbing from
_ctx_consensus_outcome, while preserving the existing FAILED and COMPLETED
outcomes and auxiliary-task waiting behavior.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

tests/unittest/others/test_cache_transceiver_precheck_run.py (2)

83-111: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add a large exact-boundary case to the parametrization.

Only the 1024 case allocates the extra reserve page. 1024 is a multiple of tokens_per_block, so num_allocated is 9 while num_prompt_blocks is 8.

7408 is not a multiple of 128. num_allocated and num_prompt_blocks are both 58, so that case would still pass if the trim in _request_block_views were removed. Add a large exact-boundary length, such as 7424 (58 blocks, 59 allocated), to protect the boundary the stacked fix targets.

♻️ Proposed parametrization change

-@pytest.mark.parametrize(("prompt_len", "expected_blocks"), ((1024, 8), (7408, 58)))
+@pytest.mark.parametrize(("prompt_len", "expected_blocks"), ((1024, 8), (7408, 58), (7424, 58)))

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/others/test_cache_transceiver_precheck_run.py` around lines 83
- 111, Extend the parametrization of
test_request_block_views_excludes_untransferred_speculative_page with a large
exact-boundary prompt length such as 7424 and expect 58 verified blocks. Keep
the existing 1024 and 7408 cases unchanged so the test covers both reserve-page
allocation and the large exact-boundary trim in _request_block_views.

---

454-464: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add coverage for the DISAGG_TRANS_ERROR state branch.

Request 1 is in the failed list here, so ctx_finish_wave raises at the failed_pairs check. The later branch that inspects req.state == self.llm_request_state.DISAGG_TRANS_ERROR is never reached by any test in this file.

Add a case where block-all reports the rid as completed but the request state is the error state. That path also must retain pages.

💚 Proposed added test

def test_ctx_finish_wave_retains_pages_on_disagg_trans_error_state(monkeypatch):
    runner, events = _ctx_finish_runner(monkeypatch, lambda _n: ([101], []))
    reqs = {0: types.SimpleNamespace(py_request_id=101, state="error")}

    with pytest.raises(rp._TransferError, match=r"ctx DISAGG_TRANS_ERROR on pairs \[0\]"):
        runner.ctx_finish_wave(reqs)

    assert events == []

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/others/test_cache_transceiver_precheck_run.py` around lines
454 - 464, Add a test alongside
test_ctx_finish_wave_retains_pages_when_request_failed that makes block-all
report the request as completed while its state is DISAGG_TRANS_ERROR, using
_ctx_finish_runner with no failed IDs and a request for pair 0. Assert
ctx_finish_wave raises the DISAGG_TRANS_ERROR-specific _TransferError and that
events remains empty, confirming pages are retained.